"""
@Author    : ghenyar
@Time      : 2025/8/28 11:09
@File      : __init__.py
@Desc      : webSocket
"""

import asyncio
import json
from typing import Dict, Union, Optional, List
from fastapi import WebSocket, WebSocketDisconnect


class WS:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, websocket: WebSocket):
        """
        第一次握手并绑定用户或管理员id
        :param websocket: websocket
        :return:
        """
        await websocket.accept()
        data = await self._processing_params(websocket)
        user_id = data.get("user_id")
        self.active_connections[user_id] = websocket

    async def heartbeat(self, websocket: WebSocket):
        """
        心跳机制
        :param websocket: websocket
        :return:
        """
        try:
            while True:
                await websocket.send_text("ping")
                await asyncio.sleep(30)  # 每30秒发送一次心跳消息
        except WebSocketDisconnect:
            await self.disconnect(websocket)

    async def handshake(self, user_id: str, websocket: WebSocket):
        # 后端握手(链接)
        self.active_connections[user_id] = websocket

    async def send(self, websocket: WebSocket):
        """
        发送消息
        :param websocket: websocket
        :return:
        """
        while True:
            data = await self._processing_params(websocket)
            if data == "pong":
                continue
            # 获取发送人id
            form_id = await self._get_from_uid(websocket)
            to_id = await self._preprocessingToUid(data.get("to_id"))
            if to_id:
                if len(to_id) == 1:
                    await self.send_personal_message(to_id=to_id[0], message=data.get("message"),
                                                     from_id=form_id)
                else:
                    await self.send_batch_users(to_ids=to_id, message=data.get("message"), from_id=form_id)

    async def send_personal_message(self, to_id: Union[str, int], message: Optional[str] = None,
                                    data: Optional[Dict] = None, from_id: Optional[int] = None):
        """
        给指定用户(一对一)发送消息
        :param to_id: 目标用户id
        :param message: 消息内容
        :param data: 其他数据
        :param from_id: 发送人  当from_id None是表示系统
        :return:
        """
        await self._send_message([to_id], message, data, from_id)

    async def send_batch_users(self, to_ids: List[Union[str, int]], message: Optional[str] = None,
                               data: Optional[Dict] = None, from_id: Optional[int] = None):
        """
        批量发送消息
        :param to_ids: 多个用户id
        :param message: 消息内容
        :param data: 其他数据
        :param from_id: 发送人  当from_id None是表示系统
        :return:
        """
        await self._send_message(to_ids, message, data, from_id)

    async def _send_message(self, to_ids: List[Union[str, int]], message: Optional[str] = None,
                            data: Optional[Dict] = None, from_id: Optional[int] = None):
        """
        发送消息的内部方法
        :param to_ids: 目标用户id列表
        :param message: 消息内容
        :param data: 其他数据
        :param from_id: 发送人  当from_id None是表示系统
        :return:
        """
        send_data = {
            "from_id": from_id,
            "message": message,
            "data": data
        }
        for to_id in to_ids:
            to_id = str(to_id) if isinstance(to_id, int) else to_id
            if to_id in self.active_connections:
                websocket = self.active_connections[to_id]
                await websocket.send_json(send_data)

    async def broadcast(self, message: Optional[str] = None,
                        data: Optional[Dict] = None):
        """
        发送广播
        :param message: 消息内容
        :param data: 其他数据
        :return:
        """
        send_data = {
            "from_id": None,
            "message": message,
            "data": data
        }
        for connection in self.active_connections.values():
            await connection.send_json(send_data)

    async def is_online(self, user_id: Union[int, str]) -> bool:
        """
        检查指定用户是否在线
        :param user_id: 用户id
        :return: True:在线  False:不在线
        """
        return str(user_id) in self.active_connections

    async def disconnect(self, websocket: WebSocket):
        """
        断开连接
        :param websocket:
        :return:
        """
        user_id = await self._get_from_uid(websocket)
        if user_id:
            del self.active_connections[user_id]
            print(f"user:{user_id} 断开websocket连接")

    async def forced_disconnect(self, user_id: str):
        # 强制断开websocket连接
        del self.active_connections[user_id]

    async def _get_from_uid(self, websocket: WebSocket) -> Optional[str]:
        """
        获取发送人id
        :param websocket:
        :return:
        """
        for uid, connection in self.active_connections.items():
            if connection == websocket:
                return uid
        return None

    @staticmethod
    async def _preprocessingToUid(to_id: Union[str, int, list]) -> Optional[List[str]]:
        """
        预处理目标用户id
        :param to_id:
        :return:
        """
        if isinstance(to_id, str):
            to_ids = [id.strip() for id in to_id.split(',') if id.strip()]
            return to_ids if to_ids else None
        elif isinstance(to_id, int):
            return [str(to_id)]
        elif isinstance(to_id, list):
            return [str(id) for id in to_id]
        return None

    @staticmethod
    async def _processing_params(websocket: WebSocket) -> Dict:
        """
        预处理websocket提交的数据
        :param websocket:
        :return:
        """
        data = await websocket.receive_text()
        return json.loads(data)


# 创建并导出全局的 wsocket 实例
wsocket = WS()
